from .data_feed import MarketDataFeed, MarketData
import json
import time

class OrderBook:
    """订单簿数据结构"""
    def __init__(self, symbol):
        self.symbol = symbol
        self.bids = {}  # 买单 {价格: 数量}
        self.asks = {}  # 卖单 {价格: 数量}
        self.last_updated = 0
        
    def update(self, bids, asks, timestamp):
        """更新订单簿数据"""
        self.bids = {price: size for price, size in bids}
        self.asks = {price: size for price, size in asks}
        self.last_updated = timestamp
    
    def get_best_bid(self):
        """获取最优买单价格"""
        if not self.bids:
            return None
        return max(self.bids.keys())
    
    def get_best_ask(self):
        """获取最优卖单价格"""
        if not self.asks:
            return None
        return min(self.asks.keys())
    
    def get_mid_price(self):
        """获取中间价"""
        best_bid = self.get_best_bid()
        best_ask = self.get_best_ask()
        if best_bid and best_ask:
            return (best_bid + best_ask) / 2
        return None
    
    def get_spread(self):
        """获取买卖价差"""
        best_bid = self.get_best_bid()
        best_ask = self.get_best_ask()
        if best_bid and best_ask:
            return best_ask - best_bid
        return None
    
    def get_order_book_imbalance(self, depth=5):
        """计算订单簿不平衡度"""
        # 按价格排序买单和卖单
        sorted_bids = sorted(self.bids.items(), key=lambda x: -x[0])[:depth]
        sorted_asks = sorted(self.asks.items(), key=lambda x: x[0])[:depth]
        
        # 计算深度内的买单和卖单总量
        bid_volume = sum(size for _, size in sorted_bids)
        ask_volume = sum(size for _, size in sorted_asks)
        
        # 计算不平衡度
        if bid_volume + ask_volume > 0:
            return (bid_volume - ask_volume) / (bid_volume + ask_volume)
        return 0

class Level2DataHandler(MarketDataFeed):
    """Level2行情数据处理器"""
    def __init__(self, config):
        super().__init__(config)
        self.order_books = {symbol: OrderBook(symbol) for symbol in config['symbols']}
        self.last_process_time = 0
        
    def _receive_raw_data(self):
        """接收Level2原始数据"""
        try:
            if self.socket:
                topic, raw_data = self.socket.recv_multipart()
                symbol = topic.decode('utf-8')
                data = json.loads(raw_data.decode('utf-8'))
                return {'symbol': symbol, 'data': data}
        except Exception as e:
            print(f"Error receiving Level2 data: {str(e)}")
        return None
    
    def _parse_data(self, raw_data):
        """解析Level2数据"""
        symbol = raw_data['symbol']
        data = raw_data['data']
        
        # 记录处理时间，用于性能监控
        current_time = time.time() * 1000  # 毫秒
        if self.last_process_time > 0:
            process_latency = current_time - self.last_process_time
            if process_latency > 5:  # 超过5ms的处理延迟
                print(f"Warning: Level2 data processing latency for {symbol}: {process_latency:.2f}ms")
        self.last_process_time = current_time
        
        # 更新订单簿
        if 'bids' in data and 'asks' in data:
            self.order_books[symbol].update(data['bids'], data['asks'], data.get('timestamp', int(current_time)))
        
        # 创建市场数据对象
        order_book = self.order_books[symbol]
        mid_price = order_book.get_mid_price()
        
        additional_fields = {
            'best_bid': order_book.get_best_bid(),
            'best_ask': order_book.get_best_ask(),
            'spread': order_book.get_spread(),
            'order_book_imbalance': order_book.get_order_book_imbalance(),
            'bids_depth': len(order_book.bids),
            'asks_depth': len(order_book.asks)
        }
        
        return MarketData(
            symbol=symbol,
            price=mid_price,
            volume=data.get('volume', 0),
            timestamp=data.get('timestamp', int(current_time)),
            exchange=self.config.get('exchange', 'unknown'),
            data_type='orderbook',
            additional_fields=additional_fields
        )
    
    def get_order_book(self, symbol):
        """获取指定标的的订单簿"""
        return self.order_books.get(symbol)
    
    def get_market_depth(self, symbol, depth=10):
        """获取市场深度数据"""
        order_book = self.order_books.get(symbol)
        if not order_book:
            return None
        
        # 获取指定深度的买卖盘
        sorted_bids = sorted(order_book.bids.items(), key=lambda x: -x[0])[:depth]
        sorted_asks = sorted(order_book.asks.items(), key=lambda x: x[0])[:depth]
        
        return {
            'bids': sorted_bids,
            'asks': sorted_asks,
            'timestamp': order_book.last_updated
        }